{-# LANGUAGE FlexibleInstances #-}

import Test.Hspec
import Test.QuickCheck
import Test.Hspec.QuickCheck

import qualified Data.Matrix as M
import Data.Matrix.SmithNormalForm
import Data.Matrix.SmithNormalForm.Internal

main :: IO ()
main = do
  hspec diagonalizeSpec
  hspec snfSpec
  hspec invariantFactorsSpec

diagonalizeSpec :: Spec
diagonalizeSpec = do
  describe "diagonalize" $ do
    prop "result is a diagonal matrix" $
      \m -> (diagonalize m :: M.Matrix Integer) `shouldSatisfy` isDiagonalMatrix

instance Arbitrary (M.Matrix Integer) where
  arbitrary = randomMatrix 

-- random (not necessarily square) matrix from 2x2 to 10x10
randomMatrix :: (Num a, Arbitrary a) => Gen (M.Matrix a)
randomMatrix = do 
  n <- choose (2,10)
  m <- choose (2,10)
  entries <- vector (n*m) 
  return (M.fromList n m entries)
                

snfSpec :: Spec
snfSpec = do
  describe "smithNormalForm" $ do
    it "is correct on sample square matrix" $ do -- from Wikipedia
      let sample = M.fromList 3 3 [2, 4, 4, -6, 6, 12, 10, 4, 16]
      smithNormalForm sample `shouldBe` M.diagonalList 3 0 [2, 2, 156]

    it "is correct on square matrix with larger entries" $ do -- src: https://www.mathworks.com/help/symbolic/smithform.html
      let largerEntryMatrix = M.fromList 5 5 [ 25, -300, 1050, -1400, 630, -300, 4800, -18900, 26880, -12600, 1050, -18900, 79380, -117600, 56700, -1400, 26880, -117600, 179200, -88200, 630, -12600, 56700, -88200, 44100]
      smithNormalForm largerEntryMatrix `shouldBe` M.diagonalList 5 0 [5, 60, 420, 840, 2520]

    it "is correct on a diagonal matrix that requires rectifying with a 0 column" $ do
      let diagonalMatrix = M.fromList 4 4 [4, 0, 0, 0, 0, 6, 0, 0, 0, 0, 8, 0, 0, 0, 0, 5]
      smithNormalForm diagonalMatrix `shouldBe` M.diagonalList 4 0 [1, 2, 4, 120]

    it "is correct on a determinant 0 matrix" $ do
      let det0example = M.fromList 4 4 [-2, 1, 1, 0, 1, -2, 1, 0, 1, 1, -3, 1, 0, 0, 1, -1]
      smithNormalForm det0example `shouldBe` M.diagonalList 4 0 [1, 1, 3, 0]

    it "is correct on a nonsquare matrix" $ do
      let nonsquare = M.fromList 2 3 [8, 4, 8, 4, 8, 4]
      smithNormalForm nonsquare `shouldBe` M.fromList 2 3 [4, 0, 0, 0, 12, 0]

    it "is correct on a nonsquare matrix with a gap of zeros" $ do
      let zerogap = M.fromList 2 5 [1, 0, -1, 0, -1, 0, 0, 1, 1, 1]
      smithNormalForm zerogap `shouldBe` M.fromList 2 5 [1, 0, 0, 0, 0, 0, 1, 0, 0, 0]

    prop "is idempotent" $
      \m -> (smithNormalForm (m :: M.Matrix Integer)) `shouldBe` (smithNormalForm . smithNormalForm) m 

invariantFactorsSpec :: Spec
invariantFactorsSpec = do
  describe "invariantFactors" $ do
    prop "divisibility conditions hold" $
      \m -> (invariantFactors (m :: M.Matrix Integer)) `shouldSatisfy` divisibilityCondition

divisibilityCondition :: [Integer] -> Bool
divisibilityCondition xs = and $ map (\(d, d') -> d `divides` d') divPairs
  where divIndices = zip [0..(length xs) - 2] [1..(length xs) - 1]
        divPairs = map (\(k, k') -> (xs !! k, xs !! k')) $ divIndices
